import os
from ament_index_python.packages import get_package_share_directory
from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument
from launch.substitutions import LaunchConfiguration, Command
from launch_ros.actions import Node
from launch_ros.actions import ComposableNodeContainer
from launch_ros.descriptions import ComposableNode
from launch_ros.parameter_descriptions import ParameterValue

def generate_launch_description():
    # ===================================================================================
    # ================================ LAUNCH ARGUMENTS =================================
    # ===================================================================================

    # --- Paths and Files ---
    urdf_file_path_arg = DeclareLaunchArgument(
        'urdf_file_path',
        default_value=os.path.join(
            get_package_share_directory('hdl_localization'), 'config', 'b2_lio_robot.urdf'
        ),
        description='Path to the robot URDF file.'
    )
    globalmap_pcd_path_arg = DeclareLaunchArgument(
        'globalmap_pcd_path',
        default_value='/home/jingzedong/B2_NAV_ws/src/hdl_Localization_ROS2/hdl_localization/data/map.pcd',
        description='Path to the globalmap PCD file.'
    )
    # --- NEW: RViz Configuration File ---
    rviz_config_arg = DeclareLaunchArgument(
        'rviz_config',
        default_value=os.path.join(
            get_package_share_directory('hdl_localization'), 'rviz', 'hdl_localization_ros2.rviz'
        ),
        description='Path to the RViz configuration file.'
    )

    # --- Basic Settings ---
    use_sim_time_arg = DeclareLaunchArgument('use_sim_time', default_value='true')

    # --- Topic Names ---
    points_topic_arg = DeclareLaunchArgument('points_topic', default_value='/rslidar_points')
    imu_topic_arg = DeclareLaunchArgument('imu_topic', default_value='/dog_imu_raw')
    odom_topic_arg = DeclareLaunchArgument('odom_topic', default_value='/odometry/imu')

    # --- Frame IDs ---
    lidar_frame_arg = DeclareLaunchArgument('lidar_frame', default_value='rslidar')
    base_link_frame_arg = DeclareLaunchArgument('base_link_frame', default_value='base_link')
    odom_frame_arg = DeclareLaunchArgument('odom_frame', default_value='odom')
    
    # --- Feature Toggles ---
    use_robot_odometry_arg = DeclareLaunchArgument(
        'use_robot_odometry', 
        default_value='false',
        description="Enable odometry-based prediction."
    )
    
    # ===================================================================================
    # ================================== NODES ==========================================
    # ===================================================================================

    # 1. Robot State Publisher
    robot_state_publisher_node = Node(
        package='robot_state_publisher',
        executable='robot_state_publisher',
        name='robot_state_publisher',
        output='screen',
        parameters=[{
            'use_sim_time': LaunchConfiguration('use_sim_time'),
            'robot_description': ParameterValue(
                Command(['cat ', LaunchConfiguration('urdf_file_path')]),
                value_type=str
            )
        }],
    )

    static_tf_publisher_node = Node(
        package='tf2_ros',
        executable='static_transform_publisher',
        name='base_link_to_robot_center',
        arguments=['0', '0', '0', '0', '0', '0', 'base_link', 'robot_center']
    )

    # 2. HDL Localization Nodes (in a component container)
    hdl_localization_container = ComposableNodeContainer(
        name='hdl_localization_container',
        namespace='',
        package='rclcpp_components',
        executable='component_container',
        composable_node_descriptions=[
            # 2.1 Global Map Server Node
            ComposableNode(
                package='hdl_localization',
                plugin='hdl_localization::GlobalmapServerNodelet',
                name='GlobalmapServerNodelet',
                parameters=[{
                    'use_sim_time': LaunchConfiguration('use_sim_time'),
                    'globalmap_pcd': LaunchConfiguration('globalmap_pcd_path'),
                    'downsample_resolution': 0.1,
                }]),
            
            # 2.2 Main Localization Node
            ComposableNode(
                package='hdl_localization',
                plugin='hdl_localization::HdlLocalizationNodelet',
                name='HdlLocalizationNodelet',
                remappings=[
                ('/ouster/points', LaunchConfiguration('points_topic')),
                ('/gpsimu_driver/imu_data', LaunchConfiguration('imu_topic')),
                ('/odom', LaunchConfiguration('odom_topic'))
                ],
                parameters=[
                    {'use_sim_time': LaunchConfiguration('use_sim_time')},
                    {'odom_child_frame_id': LaunchConfiguration('base_link_frame')},
                    {'robot_odom_frame_id': LaunchConfiguration('odom_frame')},
                    {'use_global_localization': False},
                    {'use_imu': True},
                    {'invert_acc': False},
                    {'invert_gyro': False},
                    {'enable_robot_odometry_prediction': LaunchConfiguration('use_robot_odometry')},
                    {'reg_method': 'NDT_OMP'},
                    {'ndt_resolution': 0.5},
                    {'downsample_resolution': 0.1},
                    {'specify_init_pose': True},
                    {'init_pos_x': 0.0}, {'init_pos_y': 0.0}, {'init_pos_z': 0.0},
                    {'init_ori_w': 1.0}, {'init_ori_x': 0.0}, {'init_ori_y': 0.0}, {'init_ori_z': 0.0},
                ])
        ],
        output='screen',   
    )
    
    # --- NEW: RViz2 Node ---
    rviz2_node = Node(
        package='rviz2',
        executable='rviz2',
        name='rviz2',
        arguments=['-d', LaunchConfiguration('rviz_config')],
        output='screen'
    )

    return LaunchDescription([
        # Declare all the arguments
        urdf_file_path_arg,
        globalmap_pcd_path_arg,
        rviz_config_arg, # NEW
        use_sim_time_arg,
        points_topic_arg,
        imu_topic_arg,
        odom_topic_arg,
        lidar_frame_arg,
        base_link_frame_arg,
        odom_frame_arg,
        use_robot_odometry_arg,
        static_tf_publisher_node,

        # Add the nodes to be launched
        robot_state_publisher_node,
        hdl_localization_container,
        rviz2_node,
    ])